use crate::ast::*;
use crate::error::{CompilerError, Result};
use inkwell::types::BasicType;
use inkwell::values::BasicValueEnum;
use inkwell::AddressSpace;

use super::super::LLVMBackend;

impl<'ctx> LLVMBackend<'ctx> {
    /// Handle Rc variable declaration
    pub(in crate::llvm_backend) fn handle_rc_variable_declaration(
        &mut self,
        var_name: &str,
        var_type: &Option<Type>,
        initializer: Option<&Expression>,
    ) -> Result<()> {
        let rc_type = match var_type {
            Some(Type::Rc(rc_type)) => rc_type.clone(),
            _ => {
                return Err(CompilerError::internal(
                    "Rc variable declaration requires Rc type",
                ))
            }
        };

        // Create alloca for the Rc variable (pointer or trait handle)
        let llvm_type = self.nrc_type_to_llvm_type(var_type)?;
        let alloca = self.builder.build_alloca(llvm_type, var_name)?;

        // Store the alloca in variable map
        self.variable_map
            .insert(var_name.to_string(), alloca.into());
        self.variable_type_map
            .insert(var_name.to_string(), var_type.clone().unwrap());

        // Handle initialization if present
        if let Some(init_expr) = initializer {
            let init_value = self.generate_expression(init_expr)?;
            let init_type = self.infer_expression_type(init_expr);
            let coerced = self.coerce_value_to_type(
                init_value,
                init_type.as_ref(),
                var_type.as_ref().unwrap(),
            )?;
            self.builder.build_store(alloca, coerced)?;

            let stored_value = self.builder.build_load(
                llvm_type,
                alloca,
                &format!("{}_rc_init", var_name),
            )?;
            let rc_ptr = self.rc_block_pointer_from_value(stored_value, &rc_type)?;
            self.generate_rc_increment(rc_ptr.into())?;
        } else {
            if let Some(trait_name) = self.trait_name_from_type(&rc_type.inner_type) {
                let zero = self.ensure_trait_rc_handle_type(&trait_name).const_zero();
                self.builder.build_store(alloca, zero)?;
            } else {
                let null_ptr = llvm_type.into_pointer_type().const_null();
                self.builder.build_store(alloca, null_ptr)?;
            }
        }

        Ok(())
    }

    /// Handle Weak variable declaration
    pub(in crate::llvm_backend) fn handle_weak_variable_declaration(
        &mut self,
        var_name: &str,
        var_type: &Option<Type>,
        initializer: Option<&Expression>,
    ) -> Result<()> {
        let weak_type = match var_type {
            Some(Type::Weak(weak_type)) => weak_type.clone(),
            _ => {
                return Err(CompilerError::internal(
                    "Weak variable declaration requires Weak type",
                ))
            }
        };

        let llvm_type = self.nrc_type_to_llvm_type(var_type)?;
        let alloca = self.builder.build_alloca(llvm_type, var_name)?;

        self.variable_map
            .insert(var_name.to_string(), alloca.into());
        self.variable_type_map
            .insert(var_name.to_string(), var_type.clone().unwrap());

        if let Some(init_expr) = initializer {
            let init_value = self.generate_expression(init_expr)?;
            let init_type = self.infer_expression_type(init_expr);
            let coerced = self.coerce_value_to_type(
                init_value,
                init_type.as_ref(),
                var_type.as_ref().unwrap(),
            )?;
            self.builder.build_store(alloca, coerced)?;

            let stored_value =
                self.builder
                    .build_load(llvm_type, alloca, &format!("{}_weak_init", var_name))?;
            let weak_ptr = self.weak_block_pointer_from_value(stored_value, &weak_type)?;
            self.generate_weak_increment(weak_ptr.into())?;
        } else if let Some(trait_name) = self.trait_name_from_type(&weak_type.inner_type) {
            let zero = self.ensure_trait_weak_handle_type(&trait_name).const_zero();
            self.builder.build_store(alloca, zero)?;
        } else {
            let null_ptr = llvm_type
                .into_pointer_type()
                .const_null();
            self.builder.build_store(alloca, null_ptr)?;
        }

        Ok(())
    }

    /// Handle reference counted assignment
    pub(in crate::llvm_backend) fn handle_rc_assignment(
        &mut self,
        var_name: &str,
        value_expr: &Expression,
    ) -> Result<()> {
        let var_type = self
            .variable_type_map
            .get(var_name)
            .cloned()
            .ok_or_else(|| CompilerError::internal(&format!("Undefined variable: {}", var_name)))?;
        let rc_type = match var_type {
            Type::Rc(ref rc) => rc.clone(),
            _ => {
                return Err(CompilerError::internal(&format!(
                    "Variable '{}' is not Rc",
                    var_name
                )))
            }
        };

        let var_alloca = self
            .variable_map
            .get(var_name)
            .ok_or_else(|| CompilerError::internal(&format!("Undefined variable: {}", var_name)))?
            .clone();

        let llvm_type = self.nrc_type_to_llvm_type(&Some(Type::Rc(rc_type.clone())))?;

        if let BasicValueEnum::PointerValue(alloca_ptr) = &var_alloca {
            let old_value = self.builder.build_load(
                llvm_type,
                *alloca_ptr,
                &format!("load_{}_old_ptr", var_name),
            )?;
            let old_ptr = self.rc_block_pointer_from_value(old_value, &rc_type)?;
            self.generate_rc_decrement(old_ptr.into())?;

            let new_value = self.generate_expression(value_expr)?;
            let value_type = self.infer_expression_type(value_expr);
            let coerced = self.coerce_value_to_type(
                new_value,
                value_type.as_ref(),
                &Type::Rc(rc_type.clone()),
            )?;
            self.builder.build_store(*alloca_ptr, coerced)?;

            let stored_new = self.builder.build_load(
                llvm_type,
                *alloca_ptr,
                &format!("load_{}_new_ptr", var_name),
            )?;
            let new_ptr = self.rc_block_pointer_from_value(stored_new, &rc_type)?;
            self.generate_rc_increment(new_ptr.into())?;
        }

        Ok(())
    }

    /// Handle weak reference assignment
    pub(in crate::llvm_backend) fn handle_weak_assignment(
        &mut self,
        var_name: &str,
        value_expr: &Expression,
    ) -> Result<()> {
        let var_type = self
            .variable_type_map
            .get(var_name)
            .cloned()
            .ok_or_else(|| CompilerError::internal(&format!("Undefined variable: {}", var_name)))?;
        let weak_type = match var_type {
            Type::Weak(ref weak) => weak.clone(),
            _ => {
                return Err(CompilerError::internal(&format!(
                    "Variable '{}' is not Weak",
                    var_name
                )))
            }
        };

        let var_alloca = self
            .variable_map
            .get(var_name)
            .ok_or_else(|| CompilerError::internal(&format!("Undefined variable: {}", var_name)))?
            .clone();

        let llvm_type = self.nrc_type_to_llvm_type(&Some(Type::Weak(weak_type.clone())))?;

        if let BasicValueEnum::PointerValue(alloca_ptr) = &var_alloca {
            let old_value = self.builder.build_load(
                llvm_type,
                *alloca_ptr,
                &format!("load_{}_old_weak_ptr", var_name),
            )?;
            let old_ptr = self.weak_block_pointer_from_value(old_value, &weak_type)?;
            self.generate_weak_decrement(old_ptr.into())?;

            let new_value = self.generate_expression(value_expr)?;
            let value_type = self.infer_expression_type(value_expr);
            let coerced = self.coerce_value_to_type(
                new_value,
                value_type.as_ref(),
                &Type::Weak(weak_type.clone()),
            )?;
            self.builder.build_store(*alloca_ptr, coerced)?;

            let stored_new = self.builder.build_load(
                llvm_type,
                *alloca_ptr,
                &format!("load_{}_new_weak_ptr", var_name),
            )?;
            let new_ptr = self.weak_block_pointer_from_value(stored_new, &weak_type)?;
            self.generate_weak_increment(new_ptr.into())?;
        }

        Ok(())
    }

    /// Generate reference counted object construction: rc<T>{value}
    pub(super) fn generate_rc_expression(
        &mut self,
        rc_expr: &RcExpr,
    ) -> Result<BasicValueEnum<'ctx>> {
        let mut storage_type = (*rc_expr.inner_type).clone();
        let trait_target = self.trait_name_from_type(&storage_type);
        if trait_target.is_some() {
            let init_expr = rc_expr.initializer.as_ref().ok_or_else(|| {
                CompilerError::internal("Rc<Trait> requires initializer to determine concrete type")
            })?;
            storage_type = self
                .infer_expression_type(init_expr)
                .ok_or_else(|| CompilerError::internal("Cannot infer initializer type for Rc<Trait>"))?;
        }

        // 1. 创建引用计数块
        let rc_block_value = self.generate_rc_new(&storage_type, 0)?;
        let rc_block_ptr = match rc_block_value {
            BasicValueEnum::PointerValue(ptr) => ptr,
            _ => {
                return Err(CompilerError::internal(
                    "Rc allocation should return pointer value",
                ))
            }
        };

        // 2. 如果有初始化器，设置数据
        if let Some(initializer) = &rc_expr.initializer {
            // 生成初始化器的值
            let init_value = self.generate_expression(initializer)?;

            // 获取数据指针
            let rc_block_type = self.create_rc_block_type();
            let data_ptr_field = self.builder.build_struct_gep(
                rc_block_type,
                rc_block_ptr,
                2,
                "data_ptr_field",
            )?;
            let data_ptr = self.builder.build_load(
                self.context.i8_type().ptr_type(AddressSpace::default()),
                data_ptr_field,
                "data_ptr",
            )?;

            // 将数据指针转换回正确的类型
            let inner_llvm_type = self.nrc_type_to_llvm_type(&Some(storage_type.clone()))?;
            let typed_data_ptr = self
                .builder
                .build_bitcast(
                    data_ptr.into_pointer_value(),
                    inner_llvm_type.ptr_type(AddressSpace::default()),
                    "typed_data_ptr",
                )?
                .into_pointer_value();

            // 存储初始化值
            self.builder.build_store(typed_data_ptr, init_value)?;
        }

        if let Some(trait_name) = trait_target {
            let type_name = self.extract_type_name(&storage_type).ok_or_else(|| {
                CompilerError::internal("Cannot determine concrete type name for Rc<Trait>")
            })?;
            let vtable_ptr = self.get_or_create_trait_vtable_ptr(&trait_name, &type_name)?;
            self.build_trait_rc_value(&trait_name, rc_block_ptr, vtable_ptr)
        } else {
            Ok(rc_block_ptr.into())
        }
    }
}
